package main

import (
	"bufio"
	"flag"
	"fmt"
	"net"
	"os"
	"strings"
)

var (
	protocol = flag.String("protocol", "tcp", "连接协议: tcp, udp")
	host     = flag.String("host", "localhost", "服务器地址")
	port     = flag.Int("port", 5060, "服务器端口")
)

func main() {
	flag.Parse()

	addr := fmt.Sprintf("%s:%d", *host, *port)

	switch *protocol {
	case "tcp":
		testTCPClient(addr)
	case "udp":
		testUDPClient(addr)
	default:
		fmt.Printf("不支持的协议类型: %s\n", *protocol)
		os.Exit(1)
	}
}

func testTCPClient(addr string) {
	conn, err := net.Dial("tcp", addr)
	if err != nil {
		fmt.Printf("连接TCP服务器失败: %v\n", err)
		return
	}
	defer conn.Close()

	fmt.Printf("已连接到TCP服务器: %s\n", addr)
	fmt.Println("输入命令 (REGISTER:设备ID, HEARTBEAT:设备ID, QUERY:设备ID, LIST_DEVICES, quit):")

	reader := bufio.NewReader(os.Stdin)
	for {
		fmt.Print("> ")
		input, _ := reader.ReadString('\n')
		input = strings.TrimSpace(input)

		if input == "quit" {
			break
		}

		// 发送命令
		_, err = fmt.Fprintf(conn, "%s\n", input)
		if err != nil {
			fmt.Printf("发送命令失败: %v\n", err)
			break
		}

		// 读取响应
		response, err := bufio.NewReader(conn).ReadString('\n')
		if err != nil {
			fmt.Printf("读取响应失败: %v\n", err)
			break
		}

		fmt.Printf("服务器响应: %s", response)
	}
}

func testUDPClient(addr string) {
	udpAddr, err := net.ResolveUDPAddr("udp", addr)
	if err != nil {
		fmt.Printf("解析UDP地址失败: %v\n", err)
		return
	}

	conn, err := net.DialUDP("udp", nil, udpAddr)
	if err != nil {
		fmt.Printf("连接UDP服务器失败: %v\n", err)
		return
	}
	defer conn.Close()

	fmt.Printf("已连接到UDP服务器: %s\n", addr)
	fmt.Println("输入命令 (REGISTER:设备ID, HEARTBEAT:设备ID, QUERY:设备ID, LIST_DEVICES, quit):")

	reader := bufio.NewReader(os.Stdin)
	for {
		fmt.Print("> ")
		input, _ := reader.ReadString('\n')
		input = strings.TrimSpace(input)

		if input == "quit" {
			break
		}

		// 发送命令
		_, err = conn.Write([]byte(input))
		if err != nil {
			fmt.Printf("发送命令失败: %v\n", err)
			break
		}

		// 读取响应
		buffer := make([]byte, 1024)
		n, _, err := conn.ReadFromUDP(buffer)
		if err != nil {
			fmt.Printf("读取响应失败: %v\n", err)
			break
		}

		fmt.Printf("服务器响应: %s\n", string(buffer[:n]))
	}
}